package com.sys.ai.service.impl;

import com.sys.ai.config.PostgreClient;
import com.sys.ai.domain.AiEmbeddingDb;
import com.sys.ai.domain.AiKnowledge;
import com.sys.ai.domain.AiKnowledgeFragmentation;
import com.sys.ai.service.abstracts.DbOptionAbstract;
import com.sys.ai.vo.KnowledgeFragmentationVO;
import lombok.extern.slf4j.Slf4j;
import org.postgresql.util.PGobject;
import org.springframework.stereotype.Service;

import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author LGY
 * Create by 2025/7/1 10:43
 */
@Slf4j
@Service
public class DbOptionPostgreServiceImpl extends DbOptionAbstract {

    @Override
    public void createAiKnowledge(String schemeName, AiEmbeddingDb aiEmbeddingDb) {
        PostgreClient postgreClient = PostgreClient.getInstance(aiEmbeddingDb);
        String createTableSql = "CREATE TABLE IF NOT EXISTS " + schemeName + "(\n" +
                "id SERIAL PRIMARY KEY,\n" +
                fragmentationField + " TEXT NOT NULL,\n" +
                fragmentationIdField + " TEXT NOT NULL,\n" +
                fragmentationEmbField + " vector NOT NULL\n" +
                ")";
        log.info("创建表格SQL：{}", createTableSql);
        // 创建表格
        postgreClient.executeUpdate(createTableSql, null);
        log.info("向量数据库 - 创建表格成功 - {}", schemeName);
    }

    @Override
    public void saveAiKnowledge(List<AiKnowledgeFragmentation> aiKnowledgeFragmentationList, String schemeName, AiEmbeddingDb aiEmbeddingDb) {
        PostgreClient postgreClient = PostgreClient.getInstance(aiEmbeddingDb);
        // 插入数据
        aiKnowledgeFragmentationList.forEach(item -> {
            try {
                String insertSql = "INSERT INTO " + schemeName + "(" + fragmentationField + "," + fragmentationIdField + "," + fragmentationEmbField + ") VALUES(?,?,?)";
                String fragmentationEmbStr = Arrays.toString(item.getFragmentationEmbedding());
                PGobject fragmentationObj = new PGobject();
                fragmentationObj.setType("vector");
                fragmentationObj.setValue(fragmentationEmbStr);
                List<Object> insertParam = List.of(item.getFragmentation(), item.getId(), fragmentationObj);
                log.info("插入数据SQL：{}", insertSql);
                log.info("插入数据参数：{}", insertParam);
                postgreClient.executeUpdate(insertSql, insertParam);
            } catch (Exception e) {
                throw new RuntimeException("插入失败");
            }
        });
        log.info("向量数据库 - 插入数据成功 - {}", schemeName);
    }

    public List<KnowledgeFragmentationVO> query(Float[] questionEmbedding, AiKnowledge aiKnowledge, AiEmbeddingDb aiEmbeddingDb) {
        List<KnowledgeFragmentationVO> data = new ArrayList<>();
        try {
            PostgreClient postgreClient = PostgreClient.getInstance(aiEmbeddingDb);
            String vectorDb = aiKnowledge.getVectorDb();
            String sql = "SELECT * FROM " + vectorDb +
                    " WHERE " + fragmentationEmbField + " <-> ? < " + aiEmbeddingDb.getSimilarityValue() +
                    " LIMIT 5 ";
            String fragmentationEmbStr = Arrays.toString(questionEmbedding);
            PGobject fragmentationObj = new PGobject();
            fragmentationObj.setType("vector");
            fragmentationObj.setValue(fragmentationEmbStr);
            List<Object> params = List.of(fragmentationObj);
            ResultSet resultSet = postgreClient.executeQuery(sql, params);
            while (resultSet.next()) {
                String fragmentation = resultSet.getString(fragmentationField);
                KnowledgeFragmentationVO knowledgeFragmentationVO = new KnowledgeFragmentationVO();
                knowledgeFragmentationVO.setFragmentation(fragmentation);
                data.add(knowledgeFragmentationVO);
            }
        } catch (SQLException e) {
            log.error("查询失败", e);
        }
        return data;
    }
}
